# Volatility
#
# Authors:
# Mike Auty <mike.auty@gmail.com>
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or (at
# your option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
# General Public License for more details. 
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA 
#

import re
import sre_constants
import struct
import volatility.plugins.procdump as procdump
import volatility.win32 as win32
import volatility.obj as obj
import volatility.utils as utils
import volatility.addrspace as addrspace
import volatility.debug as debug

MAX_STRING_BYTES = 260

ver_types = {
'_IMAGE_RESOURCE_DIRECTORY' : [ 0x12, {
  'Characteristics' : [ 0x0, ['unsigned long']],
  'Timestamp' : [ 0x4, ['unsigned long']],
  'MajorVersion': [ 0x8, ['unsigned short']],
  'Minorversion': [ 0xa, ['unsigned short']],
  'NamedEntriesCount': [ 0xc, ['unsigned short']],
  'IdEntriesCount': [0xe, ['unsigned short']],
  'Entries': [0x10, ['array', lambda x: x.NamedEntriesCount + x.IdEntriesCount, ['_IMAGE_RESOURCE_DIRECTORY_ENTRY']]],
} ],
'_IMAGE_RESOURCE_DIRECTORY_ENTRY': [0x8, {
  'Name' : [ 0x0, ['unsigned long']],
  'DataOffset' : [ 0x4, ['unsigned long']],
} ],
'_IMAGE_RESOURCE_DATA_ENTRY' : [0x10, {
  'DataOffset' : [0x0, ['unsigned long']],
  'Size' : [0x4, ['unsigned long']],
  'CodePage' : [0x8, ['unsigned long']],
  'Reserved' : [0xc, ['unsigned long']],
} ],
'_IMAGE_RESOURCE_DIR_STRING_U' : [0x4, {
  'Length': [0x0, ['unsigned short']],
  'Value' : [0x2, ['array', lambda x: x.Length, ['unsigned short']]],
} ],
'_VS_VERSION_INFO' : [0x26, {
  'Length': [0x0, ['unsigned short']],
  'ValueLength': [0x2, ['unsigned short']],
  'Type': [0x4, ['unsigned short']],
  'Key': [0x6, ['array', len("VS_VERSION_INFO "), ['unsigned short']]],
  'FileInfo': [lambda x: (((x.Key.obj_offset + x.Key.size() + 3) / 4) * 4), ['_VS_FIXEDFILEINFO']],
} ],
'VerStruct' : [0x26, {
  'Length': [0x0, ['unsigned short']],
  'ValueLength': [0x2, ['unsigned short']],
  'Type': [0x4, ['unsigned short']],
  'Key': [0x6, ['array', MAX_STRING_BYTES, ['unsigned short']]],
} ],
'_VS_FIXEDFILEINFO': [0x34, {
  'Signature': [0x0, ['unsigned long']],
  'StructVer': [0x4, ['unsigned long']],
  'FileVerMS': [0x8, ['unsigned long']],
  'FileVerLS': [0xC, ['unsigned long']],
  'ProdVerMS': [0x10, ['unsigned long']],
  'ProdVerLS': [0x14, ['unsigned long']],
  'FileFlagsMask': [0x18, ['unsigned long']],
  'FileFlags': [0x1C, ['unsigned long']],
  'FileOS': [0x20, ['Enumeration', {'choices': {
    0x0: 'Unknown',
    0x10000: 'DOS',
    0x20000: 'OS/2 16-bit',
    0x30000: 'OS/2 32-bit',
    0x40000: 'Windows NT',
    0x1: 'Windows 16-bit',
    0x2: 'Presentation Manager 16-bit',
    0x3: 'Presentation Manager 32-bit',
    0x4: 'Windows 32-bit',
    0x10001: 'Windows 16-bit running on DOS',
    0x10004: 'Windows 32-bit running on DOS',
    0x20002: 'Presentation Manager running on OS/2 (16-bit)',
    0x30003: 'Presentation Manager running on OS/2 (32-bit)',
    0x40004: 'Windows NT',
                                                  }} ]],
  'FileType': [0x24, ['Enumeration', {'choices': {
    0x0: 'Unknown',
    0x1: 'Application',
    0x2: 'Dynamic Link Library',
    0x3: 'Driver',
    0x4: 'Font',
    0x5: 'Virtual Device',
    0x7: 'Static Library',
                                                  }} ]],
  'FileSubType': [0x28, ['unsigned long']],
  'FileDate': [0x2C, ['WinTimeStamp']],
} ],
}

class VerStruct(obj.CType):
    """Generic Version Structure"""

    def _determine_key(self, findend = False):
        """Determines the string value for or end location of the key"""
        if self.Key != None:
            name = None
            for n in self.Key:
                if n == None:
                    return n
                # If the letter's valid, then deal with it
                if n == 0:
                    if findend:
                        return n.obj_offset + n.size()
                    name = self.obj_vm.read(self.Key.obj_offset, n.obj_offset - self.Key.obj_offset).decode("utf16", "ignore").encode("ascii", 'backslashreplace')
                    break
            return name
        return self.Key

    def get_key(self):
        """Returns the VerStruct Name"""
        return self._determine_key()

    def offset_pad(self, offset):
        """Pads an offset to a 32-bit alignment"""
        return (((offset + 3) / 4) * 4)

    def get_children(self):
        """Returns the available children"""
        offset = self.offset_pad(self._determine_key(True))
        if self.ValueLength > 0:
            # Nasty hardcoding unicode (length*2) length in here, 
            # but what else can we do?
            return self.obj_vm.read(offset, self.ValueLength * 2)
        else:
            return self._recurse_children(offset)

    def _recurse_children(self, offset):
        """Recurses thorugh the available children"""
        while offset < self.obj_offset + self.Length:
            item = obj.Object("VerStruct", offset = offset, vm = self.obj_vm, parent = self)
            if item.Length < 1 or item.get_key() == None:
                raise StopIteration("Could not recover a key for a child at offset {0}".format(item.obj_offset))
            yield item.get_key(), item.get_children()
            offset = self.offset_pad(offset + item.Length)
        raise StopIteration("No children")

class _VS_VERSION_INFO(VerStruct):
    """Version Information"""

    def get_children(self):
        """Recurses through the children of a Version Info records"""
        offset = self.offset_pad(self.FileInfo.obj_offset + self.ValueLength)
        return self._recurse_children(offset)

class _VS_FIXEDFILEINFO(obj.CType):
    """Fixed (language and codepage independent) information"""

    def file_version(self):
        """Returns the file version"""
        return self.get_version(self.FileVerMS) + "." + self.get_version(self.FileVerLS)

    def product_version(self):
        """Returns the product version"""
        return self.get_version(self.ProdVerMS) + "." + self.get_version(self.ProdVerLS)

    def get_version(self, value):
        """Returns a version in four parts"""
        version = []
        for i in range(2):
            version = [(value >> (i * 16)) & 0xFFFF] + version
        return '.'.join([str(x) for x in version])

    def file_type(self):
        """Returns the type of the file"""
        ftype = str(self.FileType)
        choices = None
        if self.FileType == 'Driver':
            choices = {
                       0x0: 'Unknown',
                       0x1: 'Printer',
                       0x2: 'Keyboard',
                       0x3: 'Language',
                       0x4: 'Display',
                       0x5: 'Mouse',
                       0x6: 'Network',
                       0x7: 'System',
                       0x8: 'Installable',
                       0x9: 'Sound',
                       0xA: 'Comms',
                       0xB: 'Input Method',
                       0xC: 'Versioned Printer',
                       }
        elif self.FileType == 'Font':
            choices = {
                       0x1: 'Raster',
                       0x2: 'Vector',
                       0x3: 'Truetype',
                       }
        if choices != None:
            subtype = obj.Object('Enumeration', 0x28, vm = self.obj_vm, parent = self, choices = choices)
            ftype += " (" + str(subtype) + ")"

        return ftype

    def flags(self):
        """Returns the file's flags"""
        data = struct.pack('=I', self.FileFlags & self.FileFlagsMask)
        addr_space = addrspace.BufferAddressSpace(self.obj_vm.get_config(), 0, data)
        bitmap = {'Debug': 0,
                  'Prerelease': 1,
                  'Patched': 2,
                  'Private Build': 3,
                  'Info Inferred': 4,
                  'Special Build' : 5,
                 }
        return obj.Object('Flags', offset = 0, vm = addr_space, bitmap = bitmap)

    def v(self):
        """Returns the value of the structure"""
        val = ("File version    : {0}\n" +
               "Product version : {1}\n" +
               "Flags           : {2}\n" +
               "OS              : {3}\n" +
               "File Type       : {4}\n" +
               "File Date       : {5}").format(self.file_version(), self.product_version(),
                                                 self.flags(), self.FileOS, self.file_type(), self.FileDate or '')
        return val

class _IMAGE_RESOURCE_DIR_STRING_U(obj.CType):
    """Handles Unicode-esque strings in IMAGE_RESOURCE_DIRECTORY structures"""
    # This is very similar to a UNICODE object, perhaps they should be merged somehow?
    def v(self):
        """Value function for _IMAGE_RESOURCE_DIR_STRING_U"""
        try:
            length = self.Length.v()
            if length > 1024:
                length = 0
            data = self.obj_vm.read(self.Value.obj_offset, length)
            return data.decode("utf16", "ignore").encode("ascii", 'backslashreplace')
        except Exception, _e:
            return ''

class _IMAGE_RESOURCE_DIRECTORY(obj.CType):
    """Handles Directory Entries"""
    def __init__(self, theType = None, offset = None, vm = None, parent = None, *args, **kwargs):
        self.sectoffset = offset
        obj.CType.__init__(self, theType = theType, offset = offset, vm = vm, parent = parent, *args, **kwargs)

    def get_entries(self):
        """Gets a tree of the entries from the top level IRD"""
        for irde in self.Entries:
            if irde != None:
                if irde.Name & 0x80000000:
                    # Points to a Name object
                    name = obj.Object("_IMAGE_RESOURCE_DIR_STRING_U", (irde.Name & 0x7FFFFFFF) + self.sectoffset, vm = self.obj_vm, parent = irde)
                else:
                    name = int(irde.Name)
                if irde.DataOffset & 0x80000000:
                    # We're another DIRECTORY
                    retobj = obj.Object("_IMAGE_RESOURCE_DIRECTORY", (irde.DataOffset & 0x7FFFFFFF) + self.sectoffset, vm = self.obj_vm, parent = irde)
                    retobj.sectoffset = self.sectoffset
                else:
                    # We're a DATA_ENTRY
                    retobj = obj.Object("_IMAGE_RESOURCE_DATA_ENTRY", irde.DataOffset + self.sectoffset, vm = self.obj_vm, parent = irde)
                yield (name, bool(irde.DataOffset & 0x80000000), retobj)

resource_types = {
 'RT_CURSOR'       : 1,
 'RT_BITMAP'       : 2,
 'RT_ICON'         : 3,
 'RT_MENU'         : 4,
 'RT_DIALOG'       : 5,
 'RT_STRING'       : 6,
 'RT_FONTDIR'      : 7,
 'RT_FONT'         : 8,
 'RT_ACCELERATOR'  : 9,
 'RT_RCDATA'       : 10,
 'RT_MESSAGETABLE' : 11,
 'RT_GROUP_CURSOR' : 12,
 'RT_GROUP_ICON'   : 14,
 'RT_VERSION'      : 16,
 'RT_DLGINCLUDE'   : 17,
 'RT_PLUGPLAY'     : 19,
 'RT_VXD'          : 20,
 'RT_ANICURSOR'    : 21,
 'RT_ANIICON'      : 22,
 'RT_HTML'         : 23,
}

class VerInfo(procdump.ProcExeDump):
    """Prints out the version information from PE images"""

    def __init__(self, config, *args):
        procdump.ProcExeDump.__init__(self, config, *args)
        config.remove_option("OFFSET")
        config.remove_option("PID")
        config.add_option("OFFSET", short_option = "o", type = 'int',
                          help = "Offset of the module to print the version information for")
        config.add_option('REGEX', short_option = "r", default = None,
                          help = 'Dump modules matching REGEX')
        config.add_option('IGNORE-CASE', short_option = 'i', action = 'store_true',
                      help = 'ignore case in pattern match', default = False)

    def calculate(self):
        """Returns a unique list of modules"""
        addr_space = utils.load_as(self._config)
        for cls in [_IMAGE_RESOURCE_DIRECTORY, _IMAGE_RESOURCE_DIR_STRING_U, _VS_FIXEDFILEINFO, _VS_VERSION_INFO, VerStruct]:
            addr_space.profile.object_classes[cls.__name__] = cls
        addr_space.profile.add_types(ver_types)

        if self._config.REGEX is not None:
            try:
                if self._config.IGNORE_CASE:
                    module_pattern = re.compile(self._config.REGEX, flags = sre_constants.SRE_FLAG_IGNORECASE)
                else:
                    module_pattern = re.compile(self._config.REGEX)
            except sre_constants.error, e:
                debug.error('Regular expression parsing error: {0}'.format(e))

        if self._config.OFFSET is not None:
            if not addr_space.is_valid_address(self._config.OFFSET):
                debug.error("Specified offset is not valid for the provided address space")
            yield addr_space, self._config.OFFSET
            raise StopIteration

        tasks = win32.tasks.pslist(addr_space)

        for task in tasks:
            for m in self.list_modules(task):
                if self._config.REGEX is not None:
                    if not (module_pattern.search(str(m.FullDllName))
                            or module_pattern.search(str(m.BaseDllName))):
                        continue

                yield task.get_process_address_space(), m

    def get_section_name(self, section):
        """Returns the null-terminated section name"""
        array = ''.join([chr(x) for x in section.Name]) + "\x00"
        return array[:array.index("\x00")]

    def display_unicode(self, string):
        """Renders a UTF16 string"""
        if string is None:
            return ''
        return string.decode("utf16", "ignore").encode("ascii", 'backslashreplace')

    def get_version_info(self, addr_space, offset):
        """Accepts an address space and an executable image offset
        
           Returns a VS_VERSION_INFO object of NoneObject
        """
        if not addr_space.is_valid_address(offset):
            return obj.NoneObject("Disk image not resident in memory")

        nt_header = self.get_nt_header(addr_space = addr_space,
                                       base_addr = offset)
        # header = s.read(m.DllBase, nt_header.OptionalHeader.SizeOfHeaders)

        for sect in self.get_sections(addr_space, nt_header):
            if self.get_section_name(sect) == '.rsrc':
                root = obj.Object("_IMAGE_RESOURCE_DIRECTORY", offset + sect.VirtualAddress, addr_space)
                for rname, rentry, rdata in root.get_entries():
                    # We're a VERSION resource and we have subelements
                    if rname == resource_types['RT_VERSION'] and rentry:
                        for sname, sentry, sdata in rdata.get_entries():
                            # We're the single sub element of the VERSION
                            if sname == 1 and sentry:
                                # Get the string tables
                                for _stname, stentry, stdata in sdata.get_entries():
                                    if not stentry:
                                        return obj.Object("_VS_VERSION_INFO", offset = (stdata.DataOffset + offset), vm = addr_space)

    def render_text(self, outfd, data):
        """Renders the text"""
        for s, m in data:
            outfd.write(str(m.FullDllName))
            outfd.write("\n")
            vinfo = self.get_version_info(s, m.DllBase)
            if vinfo != None:
                outfd.write("  File version    : {0}\n".format(vinfo.FileInfo.file_version()))
                outfd.write("  Product version : {0}\n".format(vinfo.FileInfo.product_version()))
                outfd.write("  Flags           : {0}\n".format(vinfo.FileInfo.flags()))
                outfd.write("  OS              : {0}\n".format(vinfo.FileInfo.FileOS))
                outfd.write("  File Type       : {0}\n".format(vinfo.FileInfo.file_type()))
                outfd.write("  File Date       : {0}\n".format(vinfo.FileInfo.FileDate or ''))
                for name, children in vinfo.get_children():
                    if name == 'StringFileInfo':
                        for _codepage, strings in children:
                            for string, value in strings:
                                # Make sure value isn't a generator, and we've a subtree to deal with
                                if isinstance(value, type(strings)):
                                    outfd.write("  {0} : Subtrees not yet implemented\n".format(string))
                                else:
                                    outfd.write("  {0} : {1}\n".format(string, self.display_unicode(value)))
